import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error
import matplotlib.pyplot as plt


# 数据预处理
def load_data(file_name, sequence_length, horizon):
    data = pd.read_csv(file_name)
    data = data['Close'].values.reshape(-1, 1)

    scaler = MinMaxScaler(feature_range=(0, 1))
    data = scaler.fit_transform(data)

    X, y = [], []
    for i in range(len(data) - sequence_length - horizon + 1):
        X.append(data[i:i + sequence_length])
        y.append(data[i + sequence_length:i + sequence_length + horizon])

    X = np.array(X)
    y = np.array(y)

    print("X shape:", X.shape)  # 打印 X 的形状
    print("y shape:", y.shape)  # 打印 y 的形状

    if X.size == 0 or y.size == 0:
        raise ValueError(
            f"Generated sequences are empty. Check sequence_length ({sequence_length}) and horizon ({horizon}).")

    split = int(0.8 * len(X))
    X_train, X_test = X[:split], X[split:]
    y_train, y_test = y[:split], y[split:]

    return X_train, y_train, X_test, y_test, scaler


# 定义GRU模型
class GRUModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers, output_dim):
        super(GRUModel, self).__init__()
        self.gru = nn.GRU(input_dim, hidden_dim, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        h0 = torch.zeros(self.gru.num_layers, x.size(0), self.gru.hidden_size).to(x.device)
        out, _ = self.gru(x, h0)
        out = self.fc(out[:, -1, :])  # 取最后一个时间步的输出
        return out


# 训练和测试模型
def train_model(X_train, y_train, X_test, y_test, scaler, horizon):
    X_train = torch.tensor(X_train, dtype=torch.float32)
    y_train = torch.tensor(y_train, dtype=torch.float32).squeeze(-1)  # 去掉最后一个维度
    X_test = torch.tensor(X_test, dtype=torch.float32)
    y_test = torch.tensor(y_test, dtype=torch.float32).squeeze(-1)  # 去掉最后一个维度

    input_dim = 1
    hidden_dim = 50
    num_layers = 2
    output_dim = horizon  # 输出维度为预测的天数
    learning_rate = 0.001
    num_epochs = 100

    model = GRUModel(input_dim, hidden_dim, num_layers, output_dim)
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    for epoch in range(num_epochs):
        model.train()
        optimizer.zero_grad()
        outputs = model(X_train)
        loss = criterion(outputs, y_train)
        loss.backward()
        optimizer.step()

        if (epoch + 1) % 10 == 0:
            print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}")

    model.eval()
    with torch.no_grad():
        predictions = model(X_test)
        predictions = scaler.inverse_transform(predictions.numpy())
        y_test = scaler.inverse_transform(y_test.numpy().reshape(-1, horizon))

        # 计算每个时间步的MSE
        for i in range(horizon):
            mse = mean_squared_error(y_test[:, i], predictions[:, i])
            print(f"MSE for Day {i+1}: {mse:.4f}")

        # 绘制实际值与预测值的对比图
        plt.figure(figsize=(10, 6))
        for i in range(horizon):
            plt.plot(y_test[:, i], label=f'Actual Day {i + 1}')
            plt.plot(predictions[:, i], label=f'Predicted Day {i + 1}')
        plt.legend()
        plt.show()


# 主函数
if __name__ == "__main__":
    file_name = "stock_prices.csv"
    data = pd.read_csv(file_name)
    print("Data length:", len(data))

    sequence_length = 10  # 使用过去10天的数据
    horizon = 7  # 预测未来7天的结果
    if len(data) < sequence_length + horizon:
        raise ValueError(
            f"Data length ({len(data)}) is too short for the given sequence_length ({sequence_length}) and horizon ({horizon}).")

    X_train, y_train, X_test, y_test, scaler = load_data(file_name, sequence_length, horizon)
    train_model(X_train, y_train, X_test, y_test, scaler, horizon)